"""Base classes for data model generation.

Provides ConstraintsBase for field constraints, DataModelFieldBase for field
representation, and DataModel as the abstract base for all model types.
"""

from __future__ import annotations

import re
from abc import ABC, abstractmethod
from collections import defaultdict
from copy import deepcopy
from dataclasses import dataclass
from functools import cached_property, lru_cache
from pathlib import Path
from typing import TYPE_CHECKING, Any, ClassVar, Optional, TypeVar, Union
from warnings import warn

from jinja2 import Environment, FileSystemLoader, Template
from pydantic import Field
from typing_extensions import Self

from datamodel_code_generator.imports import (
    IMPORT_ANNOTATED,
    IMPORT_OPTIONAL,
    IMPORT_UNION,
    Import,
)
from datamodel_code_generator.reference import Reference, _BaseModel
from datamodel_code_generator.types import (
    ANY,
    NONE,
    OPTIONAL_PREFIX,
    UNION_PREFIX,
    DataType,
    Nullable,
    chain_as_tuple,
    get_optional_type,
)
from datamodel_code_generator.util import PYDANTIC_V2, ConfigDict

if TYPE_CHECKING:
    from collections.abc import Iterator

    from datamodel_code_generator import DataclassArguments

TEMPLATE_DIR: Path = Path(__file__).parents[0] / "template"

ALL_MODEL: str = "#all#"

ConstraintsBaseT = TypeVar("ConstraintsBaseT", bound="ConstraintsBase")
DataModelFieldBaseT = TypeVar("DataModelFieldBaseT", bound="DataModelFieldBase")


class ConstraintsBase(_BaseModel):
    """Base class for field constraints (min/max, patterns, etc.)."""

    unique_items: Optional[bool] = Field(None, alias="uniqueItems")  # noqa: UP045
    _exclude_fields: ClassVar[set[str]] = {"has_constraints"}
    if PYDANTIC_V2:
        model_config = ConfigDict(  # pyright: ignore[reportAssignmentType]
            arbitrary_types_allowed=True, ignored_types=(cached_property,)
        )
    else:

        class Config:
            """Pydantic v1 configuration for ConstraintsBase."""

            arbitrary_types_allowed = True
            keep_untouched = (cached_property,)

    @cached_property
    def has_constraints(self) -> bool:
        """Check if any constraint values are set."""
        return any(v is not None for v in self.dict().values())

    @staticmethod
    def merge_constraints(a: ConstraintsBaseT | None, b: ConstraintsBaseT | None) -> ConstraintsBaseT | None:
        """Merge two constraint objects, with b taking precedence over a."""
        constraints_class = None
        if isinstance(a, ConstraintsBase):  # pragma: no cover
            root_type_field_constraints = {k: v for k, v in a.dict(by_alias=True).items() if v is not None}
            constraints_class = a.__class__
        else:
            root_type_field_constraints = {}  # pragma: no cover

        if isinstance(b, ConstraintsBase):  # pragma: no cover
            model_field_constraints = {k: v for k, v in b.dict(by_alias=True).items() if v is not None}
            constraints_class = constraints_class or b.__class__
        else:
            model_field_constraints = {}

        if constraints_class is None or not issubclass(constraints_class, ConstraintsBase):  # pragma: no cover
            return None

        return constraints_class.parse_obj({
            **root_type_field_constraints,
            **model_field_constraints,
        })


@dataclass(repr=False)
class WrappedDefault:
    """Represents a default value wrapped with its type constructor."""

    value: Any
    type_name: str

    def __repr__(self) -> str:
        """Return type constructor representation, e.g., 'CountType(10)'."""
        return f"{self.type_name}({self.value!r})"


class DataModelFieldBase(_BaseModel):
    """Base class for model field representation and rendering."""

    if PYDANTIC_V2:
        model_config = ConfigDict(  # pyright: ignore[reportAssignmentType]
            arbitrary_types_allowed=True,
            defer_build=True,
        )
    else:

        class Config:
            """Pydantic v1 configuration for DataModelFieldBase."""

            arbitrary_types_allowed = True

    name: Optional[str] = None  # noqa: UP045
    default: Optional[Any] = None  # noqa: UP045
    required: bool = False
    alias: Optional[str] = None  # noqa: UP045
    data_type: DataType
    constraints: Any = None
    strip_default_none: bool = False
    nullable: Optional[bool] = None  # noqa: UP045
    parent: Optional[DataModel] = None  # noqa: UP045
    extras: dict[str, Any] = Field(default_factory=dict)
    use_annotated: bool = False
    use_serialize_as_any: bool = False
    has_default: bool = False
    use_field_description: bool = False
    use_inline_field_description: bool = False
    const: bool = False
    original_name: Optional[str] = None  # noqa: UP045
    use_default_kwarg: bool = False
    use_one_literal_as_default: bool = False
    _exclude_fields: ClassVar[set[str]] = {"parent"}
    _pass_fields: ClassVar[set[str]] = {"parent", "data_type"}
    can_have_extra_keys: ClassVar[bool] = True
    type_has_null: Optional[bool] = None  # noqa: UP045
    read_only: bool = False
    write_only: bool = False

    if not TYPE_CHECKING:
        if not PYDANTIC_V2:

            @classmethod
            def model_rebuild(
                cls,
                *,
                _types_namespace: dict[str, type] | None = None,
            ) -> None:
                """Update forward references for Pydantic v1."""
                localns = _types_namespace or {}
                cls.update_forward_refs(**localns)

        def __init__(self, **data: Any) -> None:
            """Initialize the field and set up parent relationships."""
            super().__init__(**data)
            if self.data_type.reference or self.data_type.data_types:
                self.data_type.parent = self
            self.process_const()

    def process_const(self) -> None:
        """Process const values by setting them as defaults."""
        if "const" not in self.extras:
            return
        self.default = self.extras["const"]
        self.const = True
        self.required = False
        self.nullable = False

    @property
    def type_hint(self) -> str:  # noqa: PLR0911
        """Get the type hint string for this field, including nullability."""
        type_hint = self.data_type.type_hint

        if not type_hint:
            return NONE
        if self.has_default_factory or (self.data_type.is_optional and self.data_type.type != ANY):
            return type_hint
        if self.nullable is not None:
            if self.nullable:
                return get_optional_type(type_hint, self.data_type.use_union_operator)
            return type_hint
        if self.required:
            if self.type_has_null:
                return get_optional_type(type_hint, self.data_type.use_union_operator)
            return type_hint
        if self.fall_back_to_nullable:
            return get_optional_type(type_hint, self.data_type.use_union_operator)
        return type_hint

    @property
    def imports(self) -> tuple[Import, ...]:
        """Get all imports required for this field's type hint."""
        type_hint = self.type_hint
        has_union = not self.data_type.use_union_operator and UNION_PREFIX in type_hint
        has_optional = OPTIONAL_PREFIX in type_hint
        imports: list[tuple[Import] | Iterator[Import]] = [
            iter(
                i
                for i in self.data_type.all_imports
                if not ((not has_union and i == IMPORT_UNION) or (not has_optional and i == IMPORT_OPTIONAL))
            )
        ]

        if has_optional:
            imports.append((IMPORT_OPTIONAL,))
        if self.use_annotated and self.needs_annotated_import:
            imports.append((IMPORT_ANNOTATED,))
        return chain_as_tuple(*imports)

    @property
    def docstring(self) -> str | None:
        """Get the docstring for this field from its description."""
        if self.use_field_description:
            description = self.extras.get("description", None)
            if description is not None:
                return f"{description}"
        elif self.use_inline_field_description:
            # For inline mode, only use multi-line docstring format for multi-line descriptions
            description = self.extras.get("description", None)
            if description is not None and "\n" in description:
                return f"{description}"
        return None

    @property
    def inline_field_docstring(self) -> str | None:
        """Get the inline docstring for this field if single-line."""
        if self.use_inline_field_description:
            description = self.extras.get("description", None)
            if description is not None and "\n" not in description:
                return f'"""{description}"""'
        return None

    @property
    def unresolved_types(self) -> frozenset[str]:
        """Get the set of unresolved type references."""
        return self.data_type.unresolved_types

    @property
    def field(self) -> str | None:
        """For backwards compatibility."""
        return None

    @property
    def method(self) -> str | None:
        """Get the method string for this field, if any."""
        return None

    @property
    def represented_default(self) -> str:
        """Get the repr() string of the default value."""
        return repr(self.default)

    @property
    def annotated(self) -> str | None:
        """Get the Annotated type hint content, if any."""
        return None

    @property
    def needs_annotated_import(self) -> bool:
        """Check if this field requires the Annotated import."""
        return bool(self.annotated)

    @property
    def needs_meta_import(self) -> bool:  # pragma: no cover
        """Check if this field requires the Meta import (msgspec only)."""
        return False

    @property
    def has_default_factory(self) -> bool:
        """Check if this field has a default_factory."""
        return "default_factory" in self.extras

    @property
    def fall_back_to_nullable(self) -> bool:
        """Check if optional fields should be nullable by default."""
        return True

    def copy_deep(self) -> Self:
        """Create a deep copy of this field to avoid mutating the original."""
        copied = self.copy()
        copied.parent = None
        copied.data_type = self.data_type.copy()
        if self.data_type.data_types:
            copied.data_type.data_types = [dt.copy() for dt in self.data_type.data_types]
        return copied


@lru_cache
def get_template(template_file_path: Path) -> Template:
    """Load and cache a Jinja2 template from the template directory."""
    loader = FileSystemLoader(str(TEMPLATE_DIR / template_file_path.parent))
    environment: Environment = Environment(loader=loader)  # noqa: S701
    return environment.get_template(template_file_path.name)


def sanitize_module_name(name: str, *, treat_dot_as_module: bool) -> str:
    """Sanitize a module name by replacing invalid characters."""
    pattern = r"[^0-9a-zA-Z_.]" if treat_dot_as_module else r"[^0-9a-zA-Z_]"
    sanitized = re.sub(pattern, "_", name)
    if sanitized and sanitized[0].isdigit():
        sanitized = f"_{sanitized}"
    return sanitized


def get_module_path(name: str, file_path: Path | None, *, treat_dot_as_module: bool) -> list[str]:
    """Get the module path components from a name and file path."""
    if file_path:
        sanitized_stem = sanitize_module_name(file_path.stem, treat_dot_as_module=treat_dot_as_module)
        return [
            *file_path.parts[:-1],
            sanitized_stem,
            *name.split(".")[:-1],
        ]
    return name.split(".")[:-1]


def get_module_name(name: str, file_path: Path | None, *, treat_dot_as_module: bool) -> str:
    """Get the full module name from a name and file path."""
    return ".".join(get_module_path(name, file_path, treat_dot_as_module=treat_dot_as_module))


class TemplateBase(ABC):
    """Abstract base class for template-based code generation."""

    @cached_property
    @abstractmethod
    def template_file_path(self) -> Path:
        """Get the path to the template file."""
        raise NotImplementedError

    @cached_property
    def template(self) -> Template:
        """Get the cached Jinja2 template instance."""
        return get_template(self.template_file_path)

    @abstractmethod
    def render(self) -> str:
        """Render the template to a string."""
        raise NotImplementedError

    def _render(self, *args: Any, **kwargs: Any) -> str:
        """Render the template with the given arguments."""
        return self.template.render(*args, **kwargs)

    def __str__(self) -> str:
        """Return the rendered template as a string."""
        return self.render()


class BaseClassDataType(DataType):
    """DataType subclass for base class references."""


UNDEFINED: Any = object()


class DataModel(TemplateBase, Nullable, ABC):
    """Abstract base class for all data model types.

    Handles template rendering, import collection, and model relationships.
    """

    TEMPLATE_FILE_PATH: ClassVar[str] = ""
    BASE_CLASS: ClassVar[str] = ""
    DEFAULT_IMPORTS: ClassVar[tuple[Import, ...]] = ()
    IS_ALIAS: bool = False

    def __init__(  # noqa: PLR0913
        self,
        *,
        reference: Reference,
        fields: list[DataModelFieldBase],
        decorators: list[str] | None = None,
        base_classes: list[Reference] | None = None,
        custom_base_class: str | None = None,
        custom_template_dir: Path | None = None,
        extra_template_data: defaultdict[str, dict[str, Any]] | None = None,
        methods: list[str] | None = None,
        path: Path | None = None,
        description: str | None = None,
        default: Any = UNDEFINED,
        nullable: bool = False,
        keyword_only: bool = False,
        frozen: bool = False,
        treat_dot_as_module: bool = False,
        dataclass_arguments: DataclassArguments | None = None,
    ) -> None:
        """Initialize a data model with fields, base classes, and configuration."""
        self.keyword_only = keyword_only
        self.frozen = frozen
        self.dataclass_arguments: DataclassArguments = dataclass_arguments if dataclass_arguments is not None else {}
        if not self.TEMPLATE_FILE_PATH:
            msg = "TEMPLATE_FILE_PATH is undefined"
            raise Exception(msg)  # noqa: TRY002

        self._custom_template_dir: Path | None = custom_template_dir
        self.decorators: list[str] = decorators or []
        self._additional_imports: list[Import] = []
        self.custom_base_class = custom_base_class
        if base_classes:
            self.base_classes: list[BaseClassDataType] = [BaseClassDataType(reference=b) for b in base_classes]
        else:
            self.set_base_class()

        self.file_path: Path | None = path
        self.reference: Reference = reference

        self.reference.source = self

        if extra_template_data is not None:
            # The supplied defaultdict will either create a new entry,
            # or already contain a predefined entry for this type
            self.extra_template_data = extra_template_data[self.reference.path]

            # We use the full object reference path as dictionary key, but
            # we still support `name` as key because it was used for
            # `--extra-template-data` input file and we don't want to break the
            # existing behavior.
            self.extra_template_data.update(extra_template_data[self.name])
        else:
            self.extra_template_data = defaultdict(dict)

        self.fields = self._validate_fields(fields) if fields else []

        for base_class in self.base_classes:
            if base_class.reference:
                base_class.reference.children.append(self)

        if extra_template_data is not None:
            all_model_extra_template_data = extra_template_data.get(ALL_MODEL)
            if all_model_extra_template_data:
                # The deepcopy is needed here to ensure that different models don't
                # end up inadvertently sharing state (such as "base_class_kwargs")
                self.extra_template_data.update(deepcopy(all_model_extra_template_data))

        self.methods: list[str] = methods or []

        self.description = description
        for field in self.fields:
            field.parent = self

        self._additional_imports.extend(self.DEFAULT_IMPORTS)
        self.default: Any = default
        self._nullable: bool = nullable
        self._treat_dot_as_module: bool = treat_dot_as_module

    def _validate_fields(self, fields: list[DataModelFieldBase]) -> list[DataModelFieldBase]:
        names: set[str] = set()
        unique_fields: list[DataModelFieldBase] = []
        for field in fields:
            if field.name:
                if field.name in names:
                    warn(f"Field name `{field.name}` is duplicated on {self.name}", stacklevel=2)
                    continue
                names.add(field.name)
            unique_fields.append(field)
        return unique_fields

    def iter_all_fields(self, visited: set[str] | None = None) -> Iterator[DataModelFieldBase]:
        """Yield all fields including those from base classes (parent fields first)."""
        if visited is None:
            visited = set()
        if self.reference.path in visited:  # pragma: no cover
            return
        visited.add(self.reference.path)
        for base_class in self.base_classes:
            if base_class.reference and isinstance(base_class.reference.source, DataModel):
                yield from base_class.reference.source.iter_all_fields(visited)
        yield from self.fields

    def set_base_class(self) -> None:
        """Set up the base class for this model."""
        base_class = self.custom_base_class or self.BASE_CLASS
        if not base_class:
            self.base_classes = []
            return
        base_class_import = Import.from_full_path(base_class)
        self._additional_imports.append(base_class_import)
        self.base_classes = [BaseClassDataType.from_import(base_class_import)]

    @cached_property
    def template_file_path(self) -> Path:
        """Get the path to the template file, checking custom directory first."""
        template_file_path = Path(self.TEMPLATE_FILE_PATH)
        if self._custom_template_dir is not None:
            custom_template_file_path = self._custom_template_dir / template_file_path
            if custom_template_file_path.exists():
                return custom_template_file_path
        return template_file_path

    @property
    def imports(self) -> tuple[Import, ...]:
        """Get all imports required by this model and its fields."""
        return chain_as_tuple(
            (i for f in self.fields for i in f.imports),
            self._additional_imports,
        )

    @property
    def reference_classes(self) -> frozenset[str]:
        """Get all referenced class paths used by this model."""
        return frozenset(
            {r.reference.path for r in self.base_classes if r.reference}
            | {t for f in self.fields for t in f.unresolved_types}
        )

    @property
    def name(self) -> str:
        """Get the full name of this model."""
        return self.reference.name

    @property
    def duplicate_name(self) -> str:
        """Get the duplicate name for this model if it exists."""
        return self.reference.duplicate_name or ""

    @property
    def base_class(self) -> str:
        """Get the comma-separated string of base class names."""
        return ", ".join(b.type_hint for b in self.base_classes)

    @staticmethod
    def _get_class_name(name: str) -> str:
        if "." in name:
            return name.rsplit(".", 1)[-1]
        return name

    @property
    def class_name(self) -> str:
        """Get the class name without module path."""
        return self._get_class_name(self.name)

    @class_name.setter
    def class_name(self, class_name: str) -> None:
        if "." in self.reference.name:
            self.reference.name = f"{self.reference.name.rsplit('.', 1)[0]}.{class_name}"
        else:
            self.reference.name = class_name

    @property
    def duplicate_class_name(self) -> str:
        """Get the duplicate class name without module path."""
        return self._get_class_name(self.duplicate_name)

    @property
    def module_path(self) -> list[str]:
        """Get the module path components for this model."""
        return get_module_path(self.name, self.file_path, treat_dot_as_module=self._treat_dot_as_module)

    @property
    def module_name(self) -> str:
        """Get the full module name for this model."""
        return get_module_name(self.name, self.file_path, treat_dot_as_module=self._treat_dot_as_module)

    @property
    def all_data_types(self) -> Iterator[DataType]:
        """Iterate over all data types used in this model."""
        for field in self.fields:
            yield from field.data_type.all_data_types
        yield from self.base_classes

    @property
    def is_alias(self) -> bool:
        """Whether is a type alias (i.e. not an instance of BaseModel/RootModel)."""
        return self.IS_ALIAS

    @property
    def nullable(self) -> bool:
        """Check if this model is nullable."""
        return self._nullable

    @cached_property
    def path(self) -> str:
        """Get the full reference path for this model."""
        return self.reference.path

    def set_reference_path(self, new_path: str) -> None:
        """Set reference path and clear cached path property."""
        self.reference.path = new_path
        if "path" in self.__dict__:
            del self.__dict__["path"]

    def render(self, *, class_name: str | None = None) -> str:
        """Render the model to a string using the template."""
        return self._render(
            class_name=class_name or self.class_name,
            fields=self.fields,
            decorators=self.decorators,
            base_class=self.base_class,
            methods=self.methods,
            description=self.description,
            dataclass_arguments=self.dataclass_arguments,
            **self.extra_template_data,
        )


if PYDANTIC_V2:
    _rebuild_namespace = {"Union": Union, "DataModelFieldBase": DataModelFieldBase, "DataType": DataType}
    DataType.model_rebuild(_types_namespace=_rebuild_namespace)
    BaseClassDataType.model_rebuild(_types_namespace=_rebuild_namespace)
    DataModelFieldBase.model_rebuild(_types_namespace={"DataModel": DataModel})
else:
    _rebuild_namespace = {"Union": Union, "DataModelFieldBase": DataModelFieldBase, "DataType": DataType}
    DataType.model_rebuild(_types_namespace=_rebuild_namespace)
    BaseClassDataType.model_rebuild(_types_namespace=_rebuild_namespace)
    DataModelFieldBase.model_rebuild(_types_namespace={"DataModel": DataModel})
